%Run lines 1-6 to create input variables (select lines, press F9)%
clear all
working_data = []; %Col1 = time, Col2 = df/f
event_data{1,1} = {}; %Investigation type codes Col1 = S, E Col2 = A, O
event_data{2,1} = []; %Col1 = start times, Col2 = end times, Col3 = bout length, Col4 = bout intervals
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%Create Events data structure
events.social.approach.base_time{1,1} = NaN;
events.social.approach.base_trace{1,1} = NaN;
events.social.approach.invest_time{1,1} = NaN;
events.social.approach.invest_trace{1,1} = NaN;
events.social.approach.term_time{1,1} = NaN;
events.social.approach.term_trace{1,1} = NaN;

events.social.on_site.base_time{1,1} = NaN;
events.social.on_site.base_trace{1,1} = NaN;
events.social.on_site.invest_time{1,1} = NaN;
events.social.on_site.invest_trace{1,1} = NaN;
events.social.on_site.term_time{1,1} = NaN;
events.social.on_site.term_trace{1,1} = NaN;

events.empty.approach.base_time{1,1} = NaN;
events.empty.approach.base_trace{1,1} = NaN;
events.empty.approach.invest_time{1,1} = NaN;
events.empty.approach.invest_trace{1,1} = NaN;
events.empty.approach.term_time{1,1} = NaN;
events.empty.approach.term_trace{1,1} = NaN;

events.empty.on_site.base_time{1,1} = NaN;
events.empty.on_site.base_trace{1,1} = NaN;
events.empty.on_site.invest_time{1,1} = NaN;
events.empty.on_site.invest_trace{1,1} = NaN;
events.empty.on_site.term_time{1,1} = NaN;
events.empty.on_site.term_trace{1,1} = NaN;

%Turn string investigation codes into numbers and make a new "event_data"
social = "S";
approach = "A";
tempx = string(event_data{1,1});

for i = 1:size(tempx,1)
    tempy = strcmp(tempx(i,1),social);
    tempz = strcmp(tempx(i,2),approach);
    if tempy == 1
        event_data_two(i,1) = 111; %Social investigation code
    end
    if tempy ~= 1
        event_data_two(i,1) = 222; %Empty Investigation code
    end
    if tempz == 1
        event_data_two(i,2) = 333; %Approach code
    end
    if tempz ~= 1
        event_data_two(i,2) = 444; %On-site code
    end
end

tempc = event_data{2,1};
event_data_three = horzcat(event_data_two,tempc);

clear approach social temp* i event_data_two

%Combine events that have ITIs < 1s
row_count = 0;
combo_check = 0;

for i = 1:size(event_data_three,1)
    if i == 1
       row_count = row_count + 1;
       event_data_four(row_count,:) = event_data_three(i,:);
    end
    if i ~= 1 && i ~= size(event_data_three,1)
        if event_data_three(i,6) >= 1
            combo_check = 0;
            row_count = row_count + 1;
            event_data_four(row_count,:) = event_data_three(i,:);
        end
        if event_data_three(i,6) < 1 && combo_check == 0
            combo_check = 1;
            row_count = row_count + 1;
            event_data_four(row_count,1) = event_data_three(i-1,1);
            event_data_four(row_count,2) = event_data_three(i-1,2);
            event_data_four(row_count,3) = event_data_three(i-1,3);
            event_data_four(row_count,4) = event_data_three(i,4);
            event_data_four(row_count,5) = event_data_four(row_count,4) - event_data_four(row_count,3);
            event_data_four(row_count,6) = event_data_three(i,6);
            event_data_four(row_count-1,:) = NaN;
        end
        if event_data_three(i,6) < 1 && combo_check == 1
            row_count = row_count + 1;
            event_data_four(row_count,1) = event_data_four(row_count-1,1);
            event_data_four(row_count,2) = event_data_four(row_count-1,2);
            event_data_four(row_count,3) = event_data_four(row_count-1,3);
            event_data_four(row_count,4) = event_data_three(i,4);
            event_data_four(row_count,5) = event_data_four(row_count,4) - event_data_four(row_count,3);
            event_data_four(row_count,6) = event_data_three(i,6);
            event_data_four(row_count-1,:) = NaN;
        end
    end
    if i == size(event_data_three,1)
        row_count = row_count + 1;
        event_data_four(row_count,:) = event_data_three(i,:);
    end
end

row_count = 0;

for i = 1:size(event_data_four,1)
    if event_data_four(i,1) == 111 || event_data_four(i,1) == 222
        row_count = row_count + 1;
        event_data_five(row_count,:) = event_data_four(i,:);
    end
end

for i = 2:size(event_data_five,1)
    event_data_five(i,6) = event_data_five(i,3) - event_data_five(i-1,4);
end

clear combo_check event_data_three event_data_four i skip

%Exclude event if ITI < 2.5s
row_count = 0;

for i = 1:size(event_data_five,1)
    if i == 1
        row_count = row_count + 1;
        event_data_trimmed(row_count,:) = event_data_five(i,:);
        event_data_trimmed(row_count,6) = NaN;
    end
    if i ~= 1 && event_data_five(i,6) >= 2.5
        row_count = row_count + 1;
        event_data_trimmed(row_count,:) = event_data_five(i,:);
    end
end

clear event_data_five row_count i

%Collect events
row_counter_aa = 0;
row_counter_bb = 0;
row_counter_cc = 0;
row_counter_dd = 0;

for i = 1:size(event_data_trimmed,1)
    temp_class_a = event_data_trimmed(i,1);
    temp_class_b = event_data_trimmed(i,2);
    temp_start = event_data_trimmed(i,3);
    temp_end = event_data_trimmed(i,4);
    row_counter_a = 0;
    row_counter_b = 0;
    row_counter_c = 0;
    row_counter_d = 0;
    row_counter_e = 0;
    row_counter_f = 0;
    for k = 1:size(working_data,1)
        tempx = working_data(k,1);
        if tempx >= (temp_start - 3) && tempx < temp_start && temp_class_a == 111
            row_counter_a = row_counter_a + 1;
            social_base_time(row_counter_a,1) = working_data(k,1);
            social_base_trace(row_counter_a,1) = working_data(k,2);
        end
        if tempx >= (temp_start - 3) && tempx < temp_start && temp_class_a == 222
            row_counter_b = row_counter_b + 1;
            empty_base_time(row_counter_b,1) = working_data(k,1);
            empty_base_trace(row_counter_b,1) = working_data(k,2);
        end
        if tempx >= temp_start && tempx <= temp_end && temp_class_a == 111
            row_counter_c = row_counter_c + 1;
            social_invest_time(row_counter_c,1) = working_data(k,1);
            social_invest_trace(row_counter_c,1) = working_data(k,2);
        end
        if tempx >= temp_start && tempx <= temp_end && temp_class_a == 222
            row_counter_d = row_counter_d + 1;
            empty_invest_time(row_counter_d,1) = working_data(k,1);
            empty_invest_trace(row_counter_d,1) = working_data(k,2);
        end
        if tempx >= temp_end && tempx <= temp_end + 4 && temp_class_a == 111
            row_counter_e = row_counter_e + 1;
            social_term_time(row_counter_e,1) = working_data(k,1);
            social_term_trace(row_counter_e,1) = working_data(k,2);
        end
        if tempx >= temp_end && tempx <= temp_end + 4 && temp_class_a == 222
            row_counter_f = row_counter_f + 1;
            empty_term_time(row_counter_f,1) = working_data(k,1);
            empty_term_trace(row_counter_f,1) = working_data(k,2);
        end
    end
    if temp_class_a == 111 && temp_class_b == 333
        row_counter_aa = row_counter_aa + 1;
        events.social.approach.base_time{row_counter_aa,1} = social_base_time;
        events.social.approach.base_trace{row_counter_aa,1} = social_base_trace;
        events.social.approach.invest_time{row_counter_aa,1} = social_invest_time;
        events.social.approach.invest_trace{row_counter_aa,1} = social_invest_trace;
        events.social.approach.term_time{row_counter_aa,1} = social_term_time;
        events.social.approach.term_trace{row_counter_aa,1} = social_term_trace;
        clear social_* empty_*
    end
    if temp_class_a == 111 && temp_class_b == 444
        row_counter_cc = row_counter_cc + 1;
        events.social.on_site.base_time{row_counter_cc,1} = social_base_time;
        events.social.on_site.base_trace{row_counter_cc,1} = social_base_trace;
        events.social.on_site.invest_time{row_counter_cc,1} = social_invest_time;
        events.social.on_site.invest_trace{row_counter_cc,1} = social_invest_trace;
        events.social.on_site.term_time{row_counter_cc,1} = social_term_time;
        events.social.on_site.term_trace{row_counter_cc,1} = social_term_trace;
        clear social_* empty_*
    end
    if temp_class_a == 222 && temp_class_b == 333
        row_counter_bb = row_counter_bb + 1;
        events.empty.approach.base_time{row_counter_bb,1} = empty_base_time;
        events.empty.approach.base_trace{row_counter_bb,1} = empty_base_trace;
        events.empty.approach.invest_time{row_counter_bb,1} = empty_invest_time;
        events.empty.approach.invest_trace{row_counter_bb,1} = empty_invest_trace;
        events.empty.approach.term_time{row_counter_bb,1} = empty_term_time;
        events.empty.approach.term_trace{row_counter_bb,1} = empty_term_trace;
        clear social_* empty_*
    end
    if temp_class_a == 222 && temp_class_b == 444
        row_counter_dd = row_counter_dd + 1;
        events.empty.on_site.base_time{row_counter_dd,1} = empty_base_time;
        events.empty.on_site.base_trace{row_counter_dd,1} = empty_base_trace;
        events.empty.on_site.invest_time{row_counter_dd,1} = empty_invest_time;
        events.empty.on_site.invest_trace{row_counter_dd,1} = empty_invest_trace;
        events.empty.on_site.term_time{row_counter_dd,1} = empty_term_time;
        events.empty.on_site.term_trace{row_counter_dd,1} = empty_term_trace;
        clear social_* empty_*
    end
end

clear i k row_* temp* empty_* social_*

%Collect baseline means and stds for zscore calculations
base_calc.social.approach_mean = cellfun(@mean,events.social.approach.base_trace);
base_calc.social.approach_std = cellfun(@std,events.social.approach.base_trace);
base_calc.social.on_site_mean = cellfun(@mean,events.social.on_site.base_trace);
base_calc.social.on_site_std = cellfun(@std,events.social.on_site.base_trace);

base_calc.empty.approach_mean = cellfun(@mean,events.empty.approach.base_trace);
base_calc.empty.approach_std = cellfun(@std,events.empty.approach.base_trace);
base_calc.empty.on_site_mean = cellfun(@mean,events.empty.on_site.base_trace);
base_calc.empty.on_site_std = cellfun(@std,events.empty.on_site.base_trace);

%Set up data structure for event zscores
copy_events.social.approach.base_time(1:max(cellfun('size',events.social.approach.base_trace,1)),1:size(events.social.approach.base_trace,1)) = NaN;
copy_events.social.approach.base_trace_z(1:max(cellfun('size',events.social.approach.base_trace,1)),1:size(events.social.approach.base_trace,1)) = NaN;
copy_events.social.approach.invest_time(1:max(cellfun('size',events.social.approach.invest_trace,1)),1:size(events.social.approach.invest_trace,1)) = NaN;
copy_events.social.approach.invest_trace_z(1:max(cellfun('size',events.social.approach.invest_trace,1)),1:size(events.social.approach.invest_trace,1)) = NaN;
copy_events.social.approach.term_time(1:max(cellfun('size',events.social.approach.term_trace,1)),1:size(events.social.approach.term_trace,1)) = NaN;
copy_events.social.approach.term_trace_z(1:max(cellfun('size',events.social.approach.term_trace,1)),1:size(events.social.approach.term_trace,1)) = NaN;

copy_events.social.on_site.base_time(1:max(cellfun('size',events.social.on_site.base_trace,1)),1:size(events.social.on_site.base_trace,1)) = NaN;
copy_events.social.on_site.base_trace_z(1:max(cellfun('size',events.social.on_site.base_trace,1)),1:size(events.social.on_site.base_trace,1)) = NaN;
copy_events.social.on_site.invest_time(1:max(cellfun('size',events.social.on_site.invest_trace,1)),1:size(events.social.on_site.invest_trace,1)) = NaN;
copy_events.social.on_site.invest_trace_z(1:max(cellfun('size',events.social.on_site.invest_trace,1)),1:size(events.social.on_site.invest_trace,1)) = NaN;
copy_events.social.on_site.term_time(1:max(cellfun('size',events.social.on_site.term_trace,1)),1:size(events.social.on_site.term_trace,1)) = NaN;
copy_events.social.on_site.term_trace_z(1:max(cellfun('size',events.social.on_site.term_trace,1)),1:size(events.social.on_site.term_trace,1)) = NaN;

copy_events.empty.approach.base_time(1:max(cellfun('size',events.empty.approach.base_trace,1)),1:size(events.empty.approach.base_trace,1)) = NaN;
copy_events.empty.approach.base_trace_z(1:max(cellfun('size',events.empty.approach.base_trace,1)),1:size(events.empty.approach.base_trace,1)) = NaN;
copy_events.empty.approach.invest_time(1:max(cellfun('size',events.empty.approach.invest_trace,1)),1:size(events.empty.approach.invest_trace,1)) = NaN;
copy_events.empty.approach.invest_trace_z(1:max(cellfun('size',events.empty.approach.invest_trace,1)),1:size(events.empty.approach.invest_trace,1)) = NaN;
copy_events.empty.approach.term_time(1:max(cellfun('size',events.empty.approach.term_trace,1)),1:size(events.empty.approach.term_trace,1)) = NaN;
copy_events.empty.approach.term_trace_z(1:max(cellfun('size',events.empty.approach.term_trace,1)),1:size(events.empty.approach.term_trace,1)) = NaN;

copy_events.empty.on_site.base_time(1:max(cellfun('size',events.empty.on_site.base_trace,1)),1:size(events.empty.on_site.base_trace,1)) = NaN;
copy_events.empty.on_site.base_trace_z(1:max(cellfun('size',events.empty.on_site.base_trace,1)),1:size(events.empty.on_site.base_trace,1)) = NaN;
copy_events.empty.on_site.invest_time(1:max(cellfun('size',events.empty.on_site.invest_trace,1)),1:size(events.empty.on_site.invest_trace,1)) = NaN;
copy_events.empty.on_site.invest_trace_z(1:max(cellfun('size',events.empty.on_site.invest_trace,1)),1:size(events.empty.on_site.invest_trace,1)) = NaN;
copy_events.empty.on_site.term_time(1:max(cellfun('size',events.empty.on_site.term_trace,1)),1:size(events.empty.on_site.term_trace,1)) = NaN;
copy_events.empty.on_site.term_trace_z(1:max(cellfun('size',events.empty.on_site.term_trace,1)),1:size(events.empty.on_site.term_trace,1)) = NaN;

%Collect Event Times
for i = 1:size(events.social.approach.base_time,1)
    tempx = events.social.approach.base_time{i,1};
    tempy = events.social.approach.invest_time{i,1};
    tempz = events.social.approach.term_time{i,1};
    tempa = tempx - tempy(1,1);
    tempb = tempy - tempy(1,1);
    tempc = tempz - tempz(1,1);
    for k = 1:size(tempa,1)
        copy_events.social.approach.base_time(k,i) = tempa(k,1);
    end
    for q = 1:size(tempb,1)
        copy_events.social.approach.invest_time(q,i) = tempb(q,1);
    end
    for m = 1:size(tempc,1)
        copy_events.social.approach.term_time(m,i) = tempc(m,1);
    end
end
clear temp* i k q m

for i = 1:size(events.social.on_site.base_time,1)
    tempx = events.social.on_site.base_time{i,1};
    tempy = events.social.on_site.invest_time{i,1};
    tempz = events.social.on_site.term_time{i,1};
    tempa = tempx - tempy(1,1);
    tempb = tempy - tempy(1,1);
    tempc = tempz - tempz(1,1);
    for k = 1:size(tempa,1)
        copy_events.social.on_site.base_time(k,i) = tempa(k,1);
    end
    for q = 1:size(tempb,1)
        copy_events.social.on_site.invest_time(q,i) = tempb(q,1);
    end
    for m = 1:size(tempc,1)
        copy_events.social.on_site.term_time(m,i) = tempc(m,1);
    end
end
clear temp* i k q m

for i = 1:size(events.empty.approach.base_time,1)
    tempx = events.empty.approach.base_time{i,1};
    tempy = events.empty.approach.invest_time{i,1};
    tempz = events.empty.approach.term_time{i,1};
    tempa = tempx - tempy(1,1);
    tempb = tempy - tempy(1,1);
    tempc = tempz - tempz(1,1);
    for k = 1:size(tempa,1)
        copy_events.empty.approach.base_time(k,i) = tempa(k,1);
    end
    for q = 1:size(tempb,1)
        copy_events.empty.approach.invest_time(q,i) = tempb(q,1);
    end
    for m = 1:size(tempc,1)
        copy_events.empty.approach.term_time(m,i) = tempc(m,1);
    end
end
clear temp* i k q m

for i = 1:size(events.empty.on_site.base_time,1)
    tempx = events.empty.on_site.base_time{i,1};
    tempy = events.empty.on_site.invest_time{i,1};
    tempz = events.empty.on_site.term_time{i,1};
    tempa = tempx - tempy(1,1);
    tempb = tempy - tempy(1,1);
    tempc = tempz - tempz(1,1);
    for k = 1:size(tempa,1)
        copy_events.empty.on_site.base_time(k,i) = tempa(k,1);
    end
    for q = 1:size(tempb,1)
        copy_events.empty.on_site.invest_time(q,i) = tempb(q,1);
    end
    for m = 1:size(tempc,1)
        copy_events.empty.on_site.term_time(m,i) = tempc(m,1);
    end
end
clear temp* i k q m

%Calculate zscores for events
for i = 1:size(events.social.approach.base_trace,1)
    tempx = events.social.approach.base_trace{i,1};
    tempy = events.social.approach.invest_trace{i,1};
    tempz = events.social.approach.term_trace{i,1};
    for k = 1:size(tempx,1)
        tempa = (tempx(k,1) - base_calc.social.approach_mean(i,1))/base_calc.social.approach_std(i,1);
        copy_events.social.approach.base_trace_z(k,i) = tempa;
    end
    for k = 1:size(tempy,1)
        tempb = (tempy(k,1) - base_calc.social.approach_mean(i,1))/base_calc.social.approach_std(i,1);
        copy_events.social.approach.invest_trace_z(k,i) = tempb;
    end
    for k = 1:size(tempz,1)
        tempc = (tempz(k,1) - base_calc.social.approach_mean(i,1))/base_calc.social.approach_std(i,1);
        copy_events.social.approach.term_trace_z(k,i) = tempc;
    end
end
clear temp*

for i = 1:size(events.social.on_site.base_trace,1)
    tempx = events.social.on_site.base_trace{i,1};
    tempy = events.social.on_site.invest_trace{i,1};
    tempz = events.social.on_site.term_trace{i,1};
    for k = 1:size(tempx,1)
        tempa = (tempx(k,1) - base_calc.social.on_site_mean(i,1))/base_calc.social.on_site_std(i,1);
        copy_events.social.on_site.base_trace_z(k,i) = tempa;
    end
    for k = 1:size(tempy,1)
        tempb = (tempy(k,1) - base_calc.social.on_site_mean(i,1))/base_calc.social.on_site_std(i,1);
        copy_events.social.on_site.invest_trace_z(k,i) = tempb;
    end
    for k = 1:size(tempz,1)
        tempc = (tempz(k,1) - base_calc.social.on_site_mean(i,1))/base_calc.social.on_site_std(i,1);
        copy_events.social.on_site.term_trace_z(k,i) = tempc;
    end
end
clear temp*

for i = 1:size(events.empty.approach.base_trace,1)
    tempx = events.empty.approach.base_trace{i,1};
    tempy = events.empty.approach.invest_trace{i,1};
    tempz = events.empty.approach.term_trace{i,1};
    for k = 1:size(tempx,1)
        tempa = (tempx(k,1) - base_calc.empty.approach_mean(i,1))/base_calc.empty.approach_std(i,1);
        copy_events.empty.approach.base_trace_z(k,i) = tempa;
    end
    for k = 1:size(tempy,1)
        tempb = (tempy(k,1) - base_calc.empty.approach_mean(i,1))/base_calc.empty.approach_std(i,1);
        copy_events.empty.approach.invest_trace_z(k,i) = tempb;
    end
    for k = 1:size(tempz,1)
        tempc = (tempz(k,1) - base_calc.empty.approach_mean(i,1))/base_calc.empty.approach_std(i,1);
        copy_events.empty.approach.term_trace_z(k,i) = tempc;
    end
end
clear temp*

for i = 1:size(events.empty.on_site.base_trace,1)
    tempx = events.empty.on_site.base_trace{i,1};
    tempy = events.empty.on_site.invest_trace{i,1};
    tempz = events.empty.on_site.term_trace{i,1};
    for k = 1:size(tempx,1)
        tempa = (tempx(k,1) - base_calc.empty.on_site_mean(i,1))/base_calc.empty.on_site_std(i,1);
        copy_events.empty.on_site.base_trace_z(k,i) = tempa;
    end
    for k = 1:size(tempy,1)
        tempb = (tempy(k,1) - base_calc.empty.on_site_mean(i,1))/base_calc.empty.on_site_std(i,1);
        copy_events.empty.on_site.invest_trace_z(k,i) = tempb;
    end
    for k = 1:size(tempz,1)
        tempc = (tempz(k,1) - base_calc.empty.on_site_mean(i,1))/base_calc.empty.on_site_std(i,1);
        copy_events.empty.on_site.term_trace_z(k,i) = tempc;
    end
end
clear temp* i k